#include "aes.h"
#include <cstring>
#include <new>

OPEN_JLIB_NS


bool AES::setup (const void* key, s32 version) noexcept {
  if (version != 128 && version != 192 && version != 256)
    return false;
  memcpy(_key_buf, key, (uwl)version >> 3);
  _key_len = version;
  return true;
}


void AES::reset () noexcept {
  memset(_key_buf, 0, (uwl)_key_len >> 3);
  _key_len = 0;
}


uwl AES::encrypt (void* dst, const void* src, uwl src_size, u8 init_vec[AES_BLOCK_SIZE]) noexcept {
  u8 pad_buf[AES_BLOCK_SIZE];
  auto blk_size = src_size & ~(AES_BLOCK_SIZE - 1);
  auto end_size = src_size & (AES_BLOCK_SIZE - 1);
  auto pad_size = AES_BLOCK_SIZE - end_size;
  auto dp = (u8*)dst;
  auto sp = (const u8*)src;
  auto kp = _key_buf;
  auto i = 0UL;
  for (; i < end_size; ++i)
    pad_buf[i] = sp[blk_size + i];
  for (; i < AES_BLOCK_SIZE; ++i)
    pad_buf[i] = (uint8_t)pad_size;
  struct aes_key_st aes_key;
  AES_set_encrypt_key(kp, _key_len, &aes_key);
  if (blk_size)
    AES_cbc_encrypt(sp, dp, blk_size, &aes_key, init_vec, AES_ENCRYPT);
  AES_cbc_encrypt(pad_buf, dp + blk_size, AES_BLOCK_SIZE, &aes_key, init_vec, AES_ENCRYPT);
  return blk_size + AES_BLOCK_SIZE;
}


swl AES::decrypt (void* dst, const void* src, uwl src_size, u8 init_vec[AES_BLOCK_SIZE]) noexcept {
  if (src_size % AES_BLOCK_SIZE)
    return -1;
  auto dp = (u8*)dst;
  auto sp = (const u8*)src;
  auto kp = _key_buf;
  struct aes_key_st aes_key;
  AES_set_decrypt_key(kp, _key_len, &aes_key);
  AES_cbc_encrypt(sp, dp, src_size, &aes_key, init_vec, AES_DECRYPT);
  return src_size - dp[src_size - 1];
}


}